/*
 * Implementation of checker "2.12    减少运算强度"
 * Author: Fengjuan Gao
 * TODO: negative integer (ie unaryoperator), including a%(-N), a/(-N), a*(-N), (-N)*a; 
 *       N*x;
 *       CompoundAssignOperator, including a%=N, a/=N, a*=N
 */
#include "checkers/ExpensiveOperationChecker.h"
#include "CFG/SACFG.h"

extern ReportSpace::Report external_bug_report;

class FindBinaryOperations : public ASTConsumer, public RecursiveASTVisitor<FindBinaryOperations> {
    private:
        ASTContext *ctx;
        clang::SourceManager* sm;
        ASTManager *manager;
        std::string enableAliasAnalysis;
        IgnLibPathConfig *IlpCfg;
    public :
        void HandleTranslationUnit(ASTContext &cxt, ASTManager *m, std::string EnableAliasAnalysis, IgnLibPathConfig &ilpCfg)  {
            ctx = &cxt;
            sm = &(ctx->getSourceManager());
            manager = m;
            enableAliasAnalysis = EnableAliasAnalysis;
            IlpCfg = &ilpCfg;
            TranslationUnitDecl *decl = ctx->getTranslationUnitDecl();
            TraverseDecl(decl);            
        }
        
        bool VisitBinaryOperator(BinaryOperator* BO) {
            SourceLocation SL = BO->getBeginLoc();
            std::string locstr = SL.printToString(*sm);
            //avoid checking library functions
            if (IlpCfg->isInIgnoreLibPaths(locstr)) {
                return true;
            }
            std::string  recommand = "";
            //warningflag 1 means report warning, 0 means report error
            bool warningflag = 1;
            //secondOperandSign 1 means >0, 0 means < 0 , eg a%(-2)
            bool secondOperandSign = 1;
            if (BO->getOpcode() == BinaryOperatorKind::BO_Rem) { 
                //a%(2^n) => a&(2^n-1) , when a>=0              
                recommand = "Hint: a%(2^n) can be replaced with a&(2^n-1)";
                if (DeclRefExpr* dre = dyn_cast<DeclRefExpr>(BO->getLHS()->IgnoreParenCasts()->IgnoreImpCasts())) {
                    std::string firstOperandName = getUnsignedVariableName(dre);
                    if (firstOperandName != "") {
                        //only report error and repair recommandation for unsignedvar%const
                        Expr* e = BO->getRHS()->IgnoreParenCasts()->IgnoreImpCasts();
                        if (UnaryOperator* u = dyn_cast<UnaryOperator>(e)) {
                            e = u->getSubExpr();
                            secondOperandSign = 0;
                        }
                        if (IntegerLiteral* c = dyn_cast<IntegerLiteral>(e)) {        
                            uint64_t cint = c->getValue().getZExtValue();
                            if (judge(cint)) {
                                warningflag = 0;
                                recommand = "Repair recommandation: "+firstOperandName+"&"+to_string(cint-1);
                            }
                        }
                        if (DeclRefExpr* dre = dyn_cast<DeclRefExpr>(BO->getRHS()->IgnoreParenCasts()->IgnoreImpCasts())) { 
                            //handle a%alias , eg. int t=4; a%t;
                            if (VarDecl* vd = dyn_cast<VarDecl>(dre->getDecl())) {
                                uint64_t cint = getConstantAlias(BO, vd);
                                if (cint != 0 && judge(cint)) {
                                    warningflag = 0;
                                    recommand = "Repair recommandation: "+firstOperandName+"&"+to_string(cint-1);
                                }
                            }
                        }
                    }
                }


                unsigned int lineNum = sm->getExpansionLineNumber(SL);
                unsigned int columnNum = sm->getExpansionColumnNumber(SL);
                std::string filename = sm->getFilename(SL).str();
                if (filename == "") 
                    return true;
                
                ReportSpace::Defect *def = new ReportSpace::Defect(ReportSpace::DefectName::ExpensiveOperation,
                                                   (ReportSpace::DefectType)warningflag,
                                                   filename, lineNum, columnNum,
                                                   printStmt(BO));
                def->addDesc(recommand);
                external_bug_report.addToReport(def);
            }
            
            if (BO->getOpcode() == BinaryOperatorKind::BO_RemAssign) { 
                //a%=(2^n) => a&=(2^n-1) , when a>=0              
                recommand = "Hint: a%=(2^n) can be replaced with a&=(2^n-1)";
                if (DeclRefExpr* dre = dyn_cast<DeclRefExpr>(BO->getLHS()->IgnoreParenCasts()->IgnoreImpCasts())) {
                    std::string firstOperandName = getUnsignedVariableName(dre);
                    if (firstOperandName != "") {
                        //only report error and repair recommandation for unsignedvar%const
                        Expr* e = BO->getRHS()->IgnoreParenCasts()->IgnoreImpCasts();
                        if (UnaryOperator* u = dyn_cast<UnaryOperator>(e)) {
                            e = u->getSubExpr();
                            secondOperandSign = 0;
                        }
                        if (IntegerLiteral* c = dyn_cast<IntegerLiteral>(e)) {        
                            uint64_t cint = c->getValue().getZExtValue();
                            if (judge(cint)) {
                                warningflag = 0;
                                recommand = "Repair recommandation: "+firstOperandName+"&="+to_string(cint-1);
                            }
                        }
                        if (DeclRefExpr* dre = dyn_cast<DeclRefExpr>(BO->getRHS()->IgnoreParenCasts()->IgnoreImpCasts())) { 
                            //handle a%alias , eg. int t=4; a%t;
                            if (VarDecl* vd = dyn_cast<VarDecl>(dre->getDecl())) {
                                uint64_t cint = getConstantAlias(BO, vd);
                                if (cint != 0 && judge(cint)) {
                                    warningflag = 0;
                                    recommand = "Repair recommandation: "+firstOperandName+"&="+to_string(cint-1);
                                }
                            }
                        }
                    }
                }

                SourceLocation SL = BO->getBeginLoc();
                unsigned int lineNum = sm->getExpansionLineNumber(SL);
                unsigned int columnNum = sm->getExpansionColumnNumber(SL);
                std::string filename = sm->getFilename(SL).str();
                if (filename == "") 
                    return true; 
                
                ReportSpace::Defect *def = new ReportSpace::Defect(ReportSpace::DefectName::ExpensiveOperation,
                                                   (ReportSpace::DefectType)warningflag,
                                                   filename, lineNum, columnNum,
                                                   printStmt(BO));
                def->addDesc(recommand);
                external_bug_report.addToReport(def);
            }
               
            if (BO->getOpcode() == BinaryOperatorKind::BO_Div) {      
                //x/(2^n) => x>>n, when x>=0          
                recommand = "Hint: x/(2^n) can be replaced with x>>n";
                if (DeclRefExpr* dre = dyn_cast<DeclRefExpr>(BO->getLHS()->IgnoreParenCasts()->IgnoreImpCasts())) {
                    std::string firstOperandName = getUnsignedVariableName(dre);
                    if (firstOperandName != "") {
                        //only report error and repair recommandation for unsignedvar/const
                        Expr* e = BO->getRHS()->IgnoreParenCasts()->IgnoreImpCasts();
                        if (UnaryOperator* u = dyn_cast<UnaryOperator>(e)) {
                            e = u->getSubExpr();
                            secondOperandSign = 0;
                        }
                        if (IntegerLiteral* c = dyn_cast<IntegerLiteral>(e)) {      
                            uint64_t cint = c->getValue().getZExtValue();
                            if (uint64_t n = judge(cint)) {
                                warningflag = 0;
                                if (secondOperandSign) { 
                                    recommand = "Repair recommandation: "+firstOperandName+">>"+to_string(n);
                                }
                                else {
                                    // x/(-2^n) => -(x>>n)
                                    recommand = "Repair recommandation: -("+firstOperandName+">>"+to_string(n)+")";
                                }
                            }
                        }
                        if (DeclRefExpr* dre = dyn_cast<DeclRefExpr>(BO->getRHS()->IgnoreParenCasts()->IgnoreImpCasts())) {  
                            if (VarDecl* vd = dyn_cast<VarDecl>(dre->getDecl())) {
                                int64_t cint = getConstantAlias(BO, vd);
                                if (cint != 0) {
                                    if (uint64_t n = judge(cint)) {
                                        warningflag = 0;
                                        if (cint>0) { 
                                            recommand = "Repair recommandation: "+firstOperandName+">>"+to_string(n);
                                        }
                                        else {
                                            // x/(-2^n) => -(x>>n)
                                            recommand = "Repair recommandation: -("+firstOperandName+">>"+to_string(n)+")";
                                        }
                                    }
                                }
                            }
                        }
                    }
                }

                SourceLocation SL = BO->getBeginLoc();
                unsigned int lineNum = sm->getExpansionLineNumber(SL);
                unsigned int columnNum = sm->getExpansionColumnNumber(SL);
                std::string filename = sm->getFilename(SL).str();
                if (filename == "") 
                    return true; 
                ReportSpace::Defect *def = new ReportSpace::Defect(ReportSpace::DefectName::ExpensiveOperation,
                                                   (ReportSpace::DefectType)warningflag,
                                                   filename, lineNum, columnNum,
                                                   printStmt(BO));
                def->addDesc(recommand);
                external_bug_report.addToReport(def);
            }           
            
            if (BO->getOpcode() == BinaryOperatorKind::BO_DivAssign) {      
                //x/=(2^n) => x>>=n, when x>=0          
                recommand = "Hint: x/=(2^n) can be replaced with x>>=n";
                if (DeclRefExpr* dre = dyn_cast<DeclRefExpr>(BO->getLHS()->IgnoreParenCasts()->IgnoreImpCasts())) {
                    std::string firstOperandName = getUnsignedVariableName(dre);
                    if (firstOperandName != "") {
                        //only report error and repair recommandation for unsignedvar/const
                        Expr* e = BO->getRHS()->IgnoreParenCasts()->IgnoreImpCasts();
                        if (UnaryOperator* u = dyn_cast<UnaryOperator>(e)) {
                            e = u->getSubExpr();
                            secondOperandSign = 0;
                        }
                        if (IntegerLiteral* c = dyn_cast<IntegerLiteral>(e)) {      
                            uint64_t cint = c->getValue().getZExtValue();
                            if (uint64_t n = judge(cint)) {
                                if (secondOperandSign) { 
                                    warningflag = 0;
                                    recommand = "Repair recommandation: "+firstOperandName+">>="+to_string(n);
                                }
                            }
                        }
                        if (DeclRefExpr* dre = dyn_cast<DeclRefExpr>(BO->getRHS()->IgnoreParenCasts()->IgnoreImpCasts())) {  
                            if (VarDecl* vd = dyn_cast<VarDecl>(dre->getDecl())) {
                                int64_t cint = getConstantAlias(BO, vd);
                                if (cint != 0) {
                                    if (uint64_t n = judge(cint)) {
                                        if (cint>0) { 
                                            warningflag = 0;
                                            recommand = "Repair recommandation: "+firstOperandName+">>="+to_string(n);
                                        }
                                    }
                                }
                            }
                        }
                    }
                }

                SourceLocation SL = BO->getBeginLoc();
                unsigned int lineNum = sm->getExpansionLineNumber(SL);
                unsigned int columnNum = sm->getExpansionColumnNumber(SL);
                std::string filename = sm->getFilename(SL).str();
                if (filename == "") 
                    return true; 
                ReportSpace::Defect *def = new ReportSpace::Defect(ReportSpace::DefectName::ExpensiveOperation,
                                                   (ReportSpace::DefectType)warningflag,
                                                   filename, lineNum, columnNum,
                                                   printStmt(BO));
                def->addDesc(recommand);
                external_bug_report.addToReport(def);
            }           
   
            if (BO->getOpcode() == BinaryOperatorKind::BO_Mul) {      
                //x*(2^n+k) => x<<n+x*k     
                recommand = "Hint: x*(2^n+k) can be replaced with x<<n+x*k";
                if (DeclRefExpr* dre = dyn_cast<DeclRefExpr>(BO->getLHS()->IgnoreParenCasts()->IgnoreImpCasts())) {
                    std::string firstOperandName = printStmt(dre);
                    if (firstOperandName != "") {
                        //only report error and repair recommandation for unsignedvar/const
                        Expr* e = BO->getRHS()->IgnoreParenCasts()->IgnoreImpCasts();
                        if (UnaryOperator* u = dyn_cast<UnaryOperator>(e)) {
                            e = u->getSubExpr();
                            secondOperandSign = 0;
                        }
                        if (IntegerLiteral* c = dyn_cast<IntegerLiteral>(e)) {      
                            uint64_t cint = c->getValue().getZExtValue();
                            std::pair<uint64_t, uint64_t> nk = getnk(cint);
                            warningflag = 0;
                            if (secondOperandSign) { 
                                // x*(2^n+k) => x<<n+x*k  
                                if (nk.second == 0) {
                                    recommand = "Repair recommandation: "+firstOperandName+"<<"+to_string(nk.first);
                                }
                                else if (nk.second == 1) {
                                    recommand = "Repair recommandation: "+firstOperandName+"<<"+to_string(nk.first)+"+"+firstOperandName;
                                }
                                else {
                                    recommand = "Repair recommandation: "+firstOperandName+"<<"+to_string(nk.first)+"+"+firstOperandName+"*"+to_string(nk.second);
                                }
                            }
                            else {
                                // x*(-2^n-k) => (-(x))*(2^n+k) => -(x)<<n-(x)*k
                                if (nk.second == 0) {
                                    recommand = "Repair recommandation: -("+firstOperandName+")<<"+to_string(nk.first);
                                }
                                else if (nk.second == 1) {
                                    recommand = "Repair recommandation: -("+firstOperandName+")<<"+to_string(nk.first)+"-("+firstOperandName+")";
                                }
                                else {                                
                                    recommand = "Repair recommandation: -("+firstOperandName+")<<"+to_string(nk.first)+"-("+firstOperandName+")*"+to_string(nk.second);
                                }
                            }                            
                        }
                        if (DeclRefExpr* dre = dyn_cast<DeclRefExpr>(BO->getRHS()->IgnoreParenCasts()->IgnoreImpCasts())) {        
                            if (VarDecl* vd = dyn_cast<VarDecl>(dre->getDecl())) {
                                int64_t cint = getConstantAlias(BO, vd);          
                                std::pair<uint64_t, uint64_t> nk = getnk(cint);                      
                                if (cint > 0) {                                  
                                    warningflag = 0;  
                                    if (nk.second == 0) {
                                        recommand = "Repair recommandation: "+firstOperandName+"<<"+to_string(nk.first);
                                    }
                                    else if (nk.second == 1) {
                                        recommand = "Repair recommandation: "+firstOperandName+"<<"+to_string(nk.first)+"+"+firstOperandName;
                                    }
                                    else {                                
                                        recommand = "Repair recommandation: "+firstOperandName+"<<"+to_string(nk.first)+"+"+firstOperandName+"*"+to_string(nk.second);
                                    }                                    
                                }
                                else if (cint < 0) {
                                    // x*(-2^n-k) => (-(x))*(2^n+k) => -(x)<<n-(x)*k
                                    warningflag = 0;
                                    if (nk.second == 0) {
                                        recommand = "Repair recommandation: -("+firstOperandName+")<<"+to_string(nk.first);
                                    }
                                    else if (nk.second == 1) {
                                        recommand = "Repair recommandation: -("+firstOperandName+")<<"+to_string(nk.first)+"-("+firstOperandName+")";
                                    }
                                    else {                                
                                        recommand = "Repair recommandation: -("+firstOperandName+")<<"+to_string(nk.first)+"-("+firstOperandName+")*"+to_string(nk.second);
                                    }                                    
                                }
                            }
                        }
                    }
                }

                SourceLocation SL = BO->getBeginLoc();
                unsigned int lineNum = sm->getExpansionLineNumber(SL);
                unsigned int columnNum = sm->getExpansionColumnNumber(SL);
                std::string filename = sm->getFilename(SL).str();
                if (filename == "") 
                    return true; 
                ReportSpace::Defect *def = new ReportSpace::Defect(ReportSpace::DefectName::ExpensiveOperation,
                                                   (ReportSpace::DefectType)warningflag,
                                                   filename, lineNum, columnNum,
                                                   printStmt(BO));
                def->addDesc(recommand);
                external_bug_report.addToReport(def);
            }              
            if (BO->getOpcode() == BinaryOperatorKind::BO_MulAssign) {      
                //x*=(2^n+k) => x<<=n+x*k     
                recommand = "Hint: x*(2^n+k) can be replaced with x<<n+x*k";
                if (DeclRefExpr* dre = dyn_cast<DeclRefExpr>(BO->getLHS()->IgnoreParenCasts()->IgnoreImpCasts())) {
                    std::string firstOperandName = printStmt(dre);
                    if (firstOperandName != "") {
                        //only report error and repair recommandation for unsignedvar/const
                        Expr* e = BO->getRHS()->IgnoreParenCasts()->IgnoreImpCasts();
                        if (UnaryOperator* u = dyn_cast<UnaryOperator>(e)) {
                            e = u->getSubExpr();
                            secondOperandSign = 0;
                        }
                        if (IntegerLiteral* c = dyn_cast<IntegerLiteral>(e)) {      
                            uint64_t cint = c->getValue().getZExtValue();
                            std::pair<uint64_t, uint64_t> nk = getnk(cint);
                            if (secondOperandSign) { 
                                // x*(2^n) => x<<n
                                if (nk.second == 0) {
                                    warningflag = 0;
                                    recommand = "Repair recommandation: "+firstOperandName+"<<="+to_string(nk.first);
                                }                                 
                            }
                          
                        }
                        if (DeclRefExpr* dre = dyn_cast<DeclRefExpr>(BO->getRHS()->IgnoreParenCasts()->IgnoreImpCasts())) {        
                            if (VarDecl* vd = dyn_cast<VarDecl>(dre->getDecl())) {
                                int64_t cint = getConstantAlias(BO, vd);                                
                                if (cint > 0) {
                                    std::pair<uint64_t, uint64_t> nk = getnk(cint);
                                    if (nk.second == 0) {
                                        warningflag = 0;
                                        recommand = "Repair recommandation: "+firstOperandName+"<<="+to_string(nk.first);
                                    }                                  
                                }
                            }
                        }
                    }
                }

                SourceLocation SL = BO->getBeginLoc();
                unsigned int lineNum = sm->getExpansionLineNumber(SL);
                unsigned int columnNum = sm->getExpansionColumnNumber(SL);
                std::string filename = sm->getFilename(SL).str();
                if (filename == "") 
                    return true; 
                ReportSpace::Defect *def = new ReportSpace::Defect(ReportSpace::DefectName::ExpensiveOperation,
                                                   (ReportSpace::DefectType)warningflag,
                                                   filename, lineNum, columnNum,
                                                   printStmt(BO));
                def->addDesc(recommand);
                external_bug_report.addToReport(def);
            }              

            return true;
        }
        
        std::string getUnsignedVariableName(DeclRefExpr* dre) {
            if (VarDecl* vd = dyn_cast<VarDecl>(dre->getDecl())) {
                if(const BuiltinType *bt = dyn_cast<BuiltinType>(vd->getType())){
                    if(bt->isUnsignedInteger()){
                        return vd->getQualifiedNameAsString();
                    }
                }
            }
            return "";
        }
        
        uint64_t judge(uint64_t v) {
            if (v&(v-1)) //not 0 means not 2^n
                return 0;
            else {
                uint64_t n = v-1;
                uint64_t cnt = 0;
                while (n) {
                    cnt++;
                    n = n>>1;
                }
                return cnt; //v==2^cnt
            }
        }
        
        std::pair<uint64_t, uint64_t> getnk(uint64_t v) {
            //v==2^n+k
            uint64_t n=0, k=0, t = v;
            while (t) {
                n = judge(t);
                if (n) { 
                    //t=2^n
                    break;
                }
                t = t-1;
                k++;
            }
            return std::make_pair(n,k);
        }
        
        std::string printStmt(Stmt* stmt) {
            std::string s;
            llvm::raw_string_ostream stream(s);
            stmt->printPretty(stream, nullptr, PrintingPolicy(LangOptions()));
            return stream.str();
        }

        int64_t getConstantAlias(BinaryOperator* BO, VarDecl* v) {
            if (enableAliasAnalysis == "false") {
                return 0;
            }
            SourceLocation SL = BO->getBeginLoc();
            unsigned int lineNum = sm->getExpansionLineNumber(SL);
            std::string filename = sm->getFilename(SL).str();
            CFGBlock* block = manager->getBlockWithLoc(filename, lineNum);
            if (!block) {
                return 0;
            }
            
            stack<const Stmt*> s;
            for(CFGBlock::iterator iter=block->begin();iter!=block->end();iter++)
            {
                CFGElement element=*(iter);
                if(element.getKind()==CFGElement::Kind::Statement)
                {
                    const Stmt* it=((CFGStmt*)&element)->getStmt();
                     SourceLocation itSL = it->getBeginLoc();
                    std::string itfilename = sm->getFilename(itSL).str();
                    unsigned int itlineNum = sm->getExpansionLineNumber(itSL);
                    if(itfilename == filename && lineNum == itlineNum)
                    {
                        break;
                    }
                    else
                    {
                        s.push(it);
                    }
                }
            }
            bool n=true;
            VarDecl* tmp = v;
            while(!s.empty())
            {
                const Stmt* it=s.top();
                s.pop();
                if (const BinaryOperator* BO = dyn_cast<BinaryOperator>(it)) {
                    if (BO->getOpcode() == BinaryOperatorKind::BO_Assign) { 
                        if (DeclRefExpr* dre = dyn_cast<DeclRefExpr>(BO->getLHS()->IgnoreParenCasts()->IgnoreImpCasts())) {
                            if (VarDecl* vd = dyn_cast<VarDecl>(dre->getDecl())) {
                                if (vd == tmp) {
                                    //match t=x for a%t 
                                    Expr* rhs = BO->getRHS()->IgnoreParenCasts()->IgnoreImpCasts();
                                    if (DeclRefExpr* tmpdre = dyn_cast<DeclRefExpr>(rhs)) {
                                        if (VarDecl* tmpvd = dyn_cast<VarDecl>(tmpdre->getDecl())) {
                                            tmp = tmpvd;
                                            continue;
                                        }
                                        else {
                                            //t=complexExpr for a%t, quit analyzing alias
                                            break;
                                        }
                                        
                                    }
                                    else {
                                        Expr* e = rhs;
                                        bool secondOperandSign = 1;
                                        if (UnaryOperator* u = dyn_cast<UnaryOperator>(e)) {
                                            e = u->getSubExpr();
                                            secondOperandSign = 0;
                                        }
                                        if (IntegerLiteral* c = dyn_cast<IntegerLiteral>(e)) {      
                                            uint64_t cint = c->getValue().getZExtValue();
                                            return secondOperandSign == 0 ? -cint : cint;
                                        }
                                        else {
                                            break;
                                        }
                                    }
                                }
                            }
                        }
                    }
                }
            }   
            return 0;         
        }
};

void ExpensiveOperationChecker::check() {
    readConfig();
    handle();
}

void ExpensiveOperationChecker::readConfig() {
    std::unordered_map<std::string, std::string> ptrConfig = configure->getOptionBlock("ExpensiveOperationChecker");
    enableAliasAnalysis = "true";
    BlockConfigsType::iterator itr = ptrConfig.find("enableAliasAnalysis");
    if (itr != ptrConfig.end()) {
        enableAliasAnalysis = itr->second;
    }
    IlpCfg = configure->parseIgnPaths("ExpensiveOperationChecker");
}

void ExpensiveOperationChecker::handle() {
    std::vector<ASTFile *> astFiles = resource->getASTFiles();
    int count = 1;
    for(ASTFile * astFile : astFiles)
    {
        process_bar(float(count++) / astFiles.size());
        ASTUnit * astUnit = manager->getASTUnit(astFile);
        ASTContext& AST = astUnit->getASTContext();
        FindBinaryOperations fbo;
        fbo.HandleTranslationUnit(AST, manager, enableAliasAnalysis, IlpCfg);
    }    
}


